package com.cupk.rotaryguard.monitor.utils;

import com.cupk.common.base.config.Contant;
import com.opencsv.CSVReader;
import com.opencsv.exceptions.CsvException;
import org.dmg.pmml.FieldName;
import org.dmg.pmml.PMML;
import org.jpmml.evaluator.*;
import org.jpmml.model.PMMLUtil;
import org.xml.sax.SAXException;

import javax.xml.bind.JAXBException;
import java.io.*;
import java.util.*;

/**
 * 分类模型
 *
 * @author tjy
 */
public class ClassificationModel {
    private Evaluator modelEvaluator;

    /**
     * 通过传入 PMML 文件路径来生成机器学习模型
     *
     * @param pmmlFileName pmml 文件路径
     */
    public ClassificationModel(String pmmlFileName) throws Exception {
        PMML pmml = null;
        try {
            if (pmmlFileName != null) {
                InputStream is = new FileInputStream(pmmlFileName);
                pmml = PMMLUtil.unmarshal(is);
                try {
                    is.close();
                } catch (IOException e) {
                    System.out.println("InputStream close error! Exception:" + e.getMessage());
                }

                ModelEvaluatorFactory modelEvaluatorFactory = ModelEvaluatorFactory.newInstance();

                this.modelEvaluator = (Evaluator) modelEvaluatorFactory.newModelEvaluator(pmml);
                modelEvaluator.verify();
            }
        } catch (SAXException e) {
            e.printStackTrace();
        } catch (JAXBException e) {
            e.printStackTrace();
        } catch (FileNotFoundException e) {
            e.printStackTrace();
        }

    }

    /**
     * 获取模型需要的特征名称
     *
     * @return
     */
    public List<String> getFeatureNames() {
        List<String> featureNames = new ArrayList<String>();

        List<InputField> inputFields = modelEvaluator.getInputFields();

        for (InputField inputField : inputFields) {
            featureNames.add(inputField.getName().toString());
        }
        return featureNames;
    }

    /**
     * 获取目标字段名称
     *
     * @return
     */
    public String getTargetName() {
        return modelEvaluator.getTargetFields().get(0).getName().toString();
    }

    /**
     * 使用模型生成概率分布
     *
     * @param arguments
     * @return
     */
    private ProbabilityDistribution getProbabilityDistribution(Map<FieldName, ?> arguments) {
        Map<FieldName, ?> evaluateResult = modelEvaluator.evaluate(arguments);

        FieldName fieldName = new FieldName(getTargetName());

        return (ProbabilityDistribution) evaluateResult.get(fieldName);

    }

    /**
     * 预测不同分类的概率
     *
     * @param arguments
     * @return
     */
    public ValueMap<String, Number> predictProba(Map<FieldName, Number> arguments) {
        ProbabilityDistribution probabilityDistribution = getProbabilityDistribution(arguments);
        return probabilityDistribution.getValues();
    }

    /**
     * 预测结果分类
     *
     * @param arguments
     * @return
     */
    public Object predict(Map<FieldName, ?> arguments) {
        ProbabilityDistribution probabilityDistribution = getProbabilityDistribution(arguments);

        return probabilityDistribution.getPrediction();
    }

    /**
     * 将csv文件读取为Map<FieldName, ?>
     *
     * @param filePath
     * @param hasHeader
     * @return
     * @throws IOException
     */
    public static Map<FieldName, ?> readCSVToMap(String filePath, boolean hasHeader) throws IOException, CsvException {
        CSVReader csvReader = new CSVReader(new FileReader(filePath));
        List<String[]> rows = csvReader.readAll();
        String[] header = null;
        if (hasHeader) {
            header = rows.get(0);
            rows.remove(0);
        }
        List<String> fieldNames = new ArrayList<String>();
        List<List<String>> data = new ArrayList<List<String>>();
        for (String[] row : rows) {
            List<String> record = new ArrayList<String>();
            for (int i = 0; i < row.length; i++) {
                if (hasHeader) {
                    fieldNames.add(header[i]);
                } else {
                    fieldNames.add("field_" + i);
                }
                record.add(row[i]);
            }
            data.add(record);
        }
        Map<FieldName, Object> arguments = new HashMap<FieldName, Object>();
        for (int i = 0; i < fieldNames.size(); i++) {
            arguments.put(new FieldName(fieldNames.get(i)), data.get(0).get(i));
        }
        return arguments;
    }


    /**
     * 本地测试方法
     *
     * @param args
     * @throws Exception
     */
    public static void main(String[] args) throws Exception {
        //模型地址
        ClassificationModel clf = new ClassificationModel("/Users/tjy/GraduationDesign/rotaryGuard/rotaryGuard-monitor/src/main/resources/model/rf_all.pmml");

        //构建待预测数据
        Map<FieldName, Number> waitPreSample = new HashMap<>();
        String[][] strings = CsvUtil.readCSV("/Users/tjy/GraduationDesign/rotaryGuard/rotaryGuard-monitor/src/main/resources/csv/feature_all.csv", false);
        int count = 0;
        for (int i = 0; i < strings.length - 2; i++) {
            for (int j = 0; j < strings[0].length - 2; j++) {
                waitPreSample.put(new FieldName(strings[0][j]), Double.parseDouble(strings[i + 1][j]));
            }
            count++;
            System.out.println(clf.predict(waitPreSample) + "  " + count);
            System.out.println(clf.predictProba(waitPreSample).toString());
            String resultPredict = null;
            if (clf.predict(waitPreSample).toString().equals("0")) {
                resultPredict = "正常";
            } else if (clf.predict(waitPreSample).toString().equals("1")) {
                resultPredict = "滚动体故障";
            } else if (clf.predict(waitPreSample).toString().equals("2")) {
                resultPredict = "外圈故障";
            } else if (clf.predict(waitPreSample).toString().equals("3")) {
                resultPredict = "内圈故障";
            }
            System.out.println(resultPredict);
        }

    }

    /**
     * 状态诊断结果
     *
     * @param featureStringMatrix
     * @return
     * @throws Exception
     */
    public static String diagnosis(String[][] featureStringMatrix) throws Exception {
        ClassificationModel clf = new ClassificationModel(Contant.RANDOM_FOREST_MODEL_PATH);
        String resultPredict = null;
        //构建待预测数据
        Map<FieldName, Number> waitPreSample = new HashMap<>();
        String[][] strings = featureStringMatrix;
        for (int i = 1; i < strings.length; i++) {
            for (int j = 0; j < strings[0].length - 2; j++) {
                waitPreSample.put(new FieldName(strings[0][j]), Double.parseDouble(strings[i][j]));
            }
            String curPredict = clf.predict(waitPreSample).toString();
            if (curPredict.equals("0")) {
                resultPredict = "正常";
            } else if (curPredict.equals("1")) {
                resultPredict = "滚动体故障";
            } else if (curPredict.equals("2")) {
                resultPredict = "外圈故障";
            } else if (curPredict.equals("3")) {
                resultPredict = "内圈故障";
            }
        }
        return resultPredict;
    }

    /**
     * 概率分布
     *
     * @param featureStringMatrix
     * @return
     * @throws Exception
     */
    public static String diagnosisPredictProb(String[][] featureStringMatrix) throws Exception {
        ClassificationModel clf = new ClassificationModel("/Users/tjy/GraduationDesign/rotaryGuard/rotaryGuard-monitor/src/main/resources/model/rf_all.pmml");
        String diagnosisPredictProb = null;
        //构建待预测数据
        Map<FieldName, Number> waitPreSample = new HashMap<>();
        String[][] strings = featureStringMatrix;
        for (int i = 1; i < strings.length; i++) {
            for (int j = 0; j < strings[0].length - 2; j++) {
                waitPreSample.put(new FieldName(strings[0][j]), Double.parseDouble(strings[i][j]));
            }
            diagnosisPredictProb = clf.predictProba(waitPreSample).toString();
        }
        return diagnosisPredictProb;
    }
}